# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import pandas as pd

class TaskSchema(object):
    def __init__(self, passage=None, question=None, options=None, label=None, answer=None, other=None):
        self.passage = passage
        self.question = question
        self.options = options
        self.label = label
        self.answer = answer
        self.other = other

    def to_dict(self):
        return {
            "passage": self.passage,
            "question": self.question,
            "options": self.options,
            "label": self.label,
            "answer": self.answer,
            "other": self.other
        }


# define README.json
class AgiInstance(object):
    def __init__(self, task_description, data_source, task_schema, output, evaluation_metric, task_example):
        self.task_description = task_description
        self.data_source = data_source
        self.task_schema = task_schema
        self.output = output
        self.evaluation_metric = evaluation_metric
        self.task_example = task_example

    def to_dict(self):
        return {
            "task description": self.task_description,
            "data source": self.data_source,
            "task schema": self.task_schema.to_dict(),
            "output": self.output,
            "evaluation metric": self.evaluation_metric,
            "task example": self.task_example
        }


class ChatGPTSchema(object):
    def __init__(self, context=None, metadata=""):
        self.context = context
        self.metadata = metadata

    def to_dict(self):
        return {
            "context": self.context,
            "metadata": self.metadata
        }


class ResultsForHumanSchema(object):
    def __init__(self,
                 index, problem_input, label,
                 model_input="", model_output="",
                 parse_result="",
                 first_stage_output="", second_stage_input="",
                 is_correct=False
                 ):
        self.index = index
        self.problem_input = problem_input
        self.model_input = model_input
        self.model_output = model_output
        self.parse_result = parse_result
        self.label = label
        self.first_stage_output = first_stage_output
        self.second_stage_input = second_stage_input
        self.is_correct = is_correct

    def to_dict(self):
        return {
            "index": self.index,
            "problem_input": self.problem_input,
            "model_input": self.model_input,
            "model_output": self.model_output,
            "parse_result": self.parse_result,
            "label": self.label,
            "is_correct": self.is_correct,
            "first_stage_output": self.first_stage_output,
            "second_stage_input": self.second_stage_input,
        }

    @staticmethod
    def to_tsv(result_list, path):
        result_json = [item.to_dict() for item in result_list]
        table = pd.json_normalize(result_json)
        table.to_excel(path, index=False)

